import torch.nn as nn
import torch
import math


class yolo_block(nn.Module):

    def __init__(self, in_channels: int, out_channels: int, conv_size: int, stride: int, padding: int, bias: bool = True, need_normal: bool = True) -> None:
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.conv = nn.Conv2d(in_channels, out_channels, conv_size, stride, padding, bias=bias)
        self.acfun = nn.LeakyReLU(0.1)
        self.need_normal = need_normal
        if need_normal:
            self.normal = nn.BatchNorm2d(out_channels)
        # 初始化权重
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = self.conv(x)
        out = self.acfun(out)
        if self.need_normal:
            out = self.normal(out)
        return out


class yolo_resnet(nn.Module):

    def __init__(self, yolo_blocks: list) -> None:
        super().__init__()
        self.in_channels = yolo_blocks[0].in_channels
        self.out_channels = yolo_blocks[-1].out_channels
        self.resnet_conv = nn.Conv2d(self.in_channels, self.out_channels, 1, 1, 0)
        self.yolo_blocks = nn.ModuleList(yolo_blocks)
        self.acfun = nn.LeakyReLU(0.1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = x
        for yolo_block in self.yolo_blocks:
            out = yolo_block(out)
        out = out + self.resnet_conv(x)
        return self.acfun(out)


class yolov1(nn.Module):
    def __init__(self, bounding_box_size=2, classes_num=20) -> None:
        super().__init__()
        self.classes_num = classes_num
        self.bounding_box_size = bounding_box_size
        self.layers = nn.Sequential(  # 448 * 448
            yolo_block(3, 64, 7, 2, 3),  # 224 * 224
            nn.MaxPool2d(2, 2),  # 112 * 112
            yolo_resnet([yolo_block(64, 192, 3, 1, 1)]),
            nn.MaxPool2d(2, 2),  # 56 * 56
            yolo_resnet([
                yolo_block(192, 128, 1, 1, 0),
                yolo_block(128, 256, 3, 1, 1),
                yolo_block(256, 256, 1, 1, 0),
                yolo_block(256, 512, 3, 1, 1),
            ]),
            nn.MaxPool2d(2, 2),  # 28 * 28
            yolo_resnet([
                yolo_block(512, 256, 1, 1, 0),
                yolo_block(256, 512, 3, 1, 1),
                yolo_block(512, 256, 1, 1, 0),
                yolo_block(256, 512, 3, 1, 1),
                yolo_block(512, 256, 1, 1, 0),
                yolo_block(256, 512, 3, 1, 1),
                yolo_block(512, 256, 1, 1, 0),
                yolo_block(256, 512, 3, 1, 1),
                yolo_block(512, 512, 1, 1, 0),
                yolo_block(512, 1024, 3, 1, 1),
            ]),
            nn.MaxPool2d(2, 2),  # 14 * 14
            yolo_resnet([
                yolo_block(1024, 512, 1, 1, 0),
                yolo_block(512, 1024, 3, 1, 1),
                yolo_block(1024, 512, 1, 1, 0),
                yolo_block(512, 1024, 3, 1, 1),
                yolo_block(1024, 1024, 3, 1, 1),
            ]),
            yolo_block(1024, 1024, 3, 2, 1),  # 7 * 7
            yolo_resnet([
                yolo_block(1024, 1024, 3, 1, 1),
                yolo_block(1024, 1024, 3, 1, 1),
            ]),
        )
        self.fc1 = nn.Linear(7 * 7 * 1024, 1024)
        self.fc2 = nn.Linear(1024, 7 * 7 * (self.bounding_box_size * 5 + self.classes_num))

        self.sigmoid = nn.Sigmoid()
        self.softmax = nn.Softmax(dim=3)

        # 初始化权重 yolo_resnet 和 yolo_block在初始化时已经初始化权重了
        for m in self.layers:
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                nn.init.constant_(m.bias, 0)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch_size = x.size(0)
        out = self.layers(x)
        out = out.view(batch_size, -1)
        out = self.fc1(out)
        out = self.fc2(out)
        out = out.view(batch_size, 7, 7, (5 * self.bounding_box_size + self.classes_num))

        bnd_coord = self.sigmoid(out[:, :, :, 0: self.bounding_box_size * 5])
        bnd_cls = self.softmax(out[:, :, :, self.bounding_box_size * 5:])
        out = torch.cat([bnd_coord, bnd_cls], dim=3)
        return out


class resnet(nn.Module):
    def __init__(self, otherModel: nn.Module, acfun: nn.Module) -> None:
        super().__init__()
        self.otherModel = otherModel
        self.acfun = acfun

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = self.otherModel(x)
        out = out + x
        if self.acfun is not None:
            out = self.acfun(out)
        return out
